import cython
import threading
import queue
import numpy as np
cimport numpy
import copy
import time
from cpython.mem cimport PyMem_Malloc, PyMem_Realloc, PyMem_Free
from asr_util import parse_hypothesis
from AsrStDecoder import warp_asr_free_stdecoder,warp_asr_alloc_stdecoder,warp_asr_asyn_decode,warp_asr_gethyps_num,warp_asr_input_online,warp_asr_input_end,warp_asr_gethyps

cdef class StreamingDecoder:
    cdef object oid
    cdef object _decoder
    cdef object hyp_num,hyps
    cdef object finish_sem
    cdef object wavque,resque

    def __cinit__(self, sid, model_path):
        self.oid = sid
        self._decoder = np.zeros((1), dtype=np.int64)
        self.hyp_num = np.zeros((1), dtype=np.int32)
        self.hyps = None
        warp_asr_alloc_stdecoder(self._decoder,model_path)
        self.wavque = queue.Queue()
        self.resque = queue.Queue()
        self.finish_sem = threading.Semaphore(0)

    def __dealloc__(self):
        warp_asr_free_stdecoder(self._decoder)

    def decode_streaming_window(self, wavbytes):
        self.wavque.put(wavbytes)
        max_num = np.zeros((1), dtype=np.int32)
        warp_asr_gethyps_num(self._decoder,max_num)
        if max_num[0] > 0:
            self.hyps = np.zeros((max_num[0]), dtype=np.int32)
            warp_asr_gethyps(self._decoder,max_num, self.hyps, self.hyp_num)
        return self.hyps
        
    def decode_streaming_finish(self):
        self.wavque.put(int(0))
        self.finish_sem.acquire()
        max_num = np.zeros((1), dtype=np.int32)
        warp_asr_gethyps_num(self._decoder,max_num)
        if max_num[0] > 0:
            self.hyps = np.zeros((max_num[0]), dtype=np.int32)
            warp_asr_gethyps(self._decoder,max_num, self.hyps, self.hyp_num)
        return self.hyps

def worker_feat(StreamingDecoder sd):
    while True:
        wavbytes = sd.wavque.get()
        hyp_num = np.zeros((1), dtype=np.int32)
        if isinstance(wavbytes, np.ndarray):
            warp_asr_input_online(sd._decoder, wavbytes, wavbytes.shape[0])    
        else:
            warp_asr_input_end(sd._decoder)
            break

def worker_decode(StreamingDecoder sd):
    warp_asr_asyn_decode(sd._decoder)
    sd.finish_sem.release()
